
import numpy as np
import os
import matplotlib.pyplot as plt
from collections import deque

data_dir = r"datalib/glue_dataset_1"

def better_label(tacdata):
    sensored_force = np.sum(tacdata[:, :, 6:9], axis=1)

    tangent_force = np.linalg.norm(sensored_force[:,0:2], axis=1)
    normal_force = -sensored_force[:, 2]
    
    # take the idx where tangent force is maximum as target index
    target_idx = np.argmax(tangent_force)
    mu1 = tangent_force[target_idx] / normal_force[target_idx]
    mu2 = tangent_force[-1] / normal_force[-1]
    mu = np.ones_like(tangent_force) * mu1

    # interpolate mu by contact center position
    contact_center_row, contact_center_col = calc_contact_center(tacdata)
    contact_direction = np.array([contact_center_row[-1] - contact_center_row[target_idx], contact_center_col[-1] - contact_center_col[target_idx]])
    contact_delta = np.stack([contact_center_row[target_idx:] - contact_center_row[target_idx], contact_center_col[target_idx:] - contact_center_col[target_idx]],axis = -1)
    print(contact_delta.shape, contact_direction.shape)
    interp_weight = contact_delta @ contact_direction / np.linalg.norm(contact_direction) / np.linalg.norm(contact_direction)
    interp_weight = np.clip(interp_weight, 0, 1)
    mu[target_idx:] = mu1 * (1 - interp_weight) + mu2 * interp_weight

    better_label = 1 - tangent_force / (normal_force * mu)
    better_label = np.clip(better_label, 0, 1)
    return better_label ,mu


def load_data(exp_dir, sensor_id):
    try:
        # print(f"Loading data from {exp_dir} for sensor {sensor_id}")
        tacdata = np.load(os.path.join(exp_dir, f"tacdata_{sensor_id}.npy"))
        true_delta = np.load(os.path.join(exp_dir, f"true_delta_{sensor_id}.npy"))
        predicted_delta = np.load(os.path.join(exp_dir, f"predicted_delta_{sensor_id}.npy"))
        better_delta,mu = better_label(tacdata)
        grasp_force = np.load(os.path.join(exp_dir, f"grasp_force_{sensor_id}.npy"))
        goal_force = np.load(os.path.join(exp_dir, f"goal_force_{sensor_id}.npy"))
    
        return tacdata, true_delta, predicted_delta, grasp_force, goal_force, better_delta ,mu
    except FileNotFoundError as e:
        print(f"File not found: {e}")
        return None, None, None, None, None, None , None

def calc_contact_center(tacdata):
    """
    Calculate the contact center from Tac3D data.
    
    Args:
        tacdata (np.ndarray): The Tac3D data array.
        
    Returns:
        np.ndarray: The contact center coordinates.
    """
    resultant_force = np.linalg.norm(tacdata[:,:, 6:9], axis=2).reshape(-1,20, 20)
    all_force = np.sum(resultant_force, axis=(1, 2))
    idx = np.arange(20)
    row = np.sum(idx[np.newaxis,:,np.newaxis] * resultant_force, axis=(1,2)) / all_force
    col = np.sum(idx[np.newaxis,np.newaxis,:] * resultant_force, axis=(1,2)) / all_force
    return row,col

class TangentForceRateCalculator:
    """Calculate the rate of change of tangent force."""
    def __init__(self, window_size=5, momentum=0.9):
        """
        Initialize the calculator with a window size.
        Args:
            window_size (int): The size of the window for calculating the rate.
        """
        self.window_size = window_size
        self.previous_tangent_force = deque(maxlen=window_size)
        self.momentum = momentum
        self.now_rate = 0.0
        x = np.arange(window_size)
        self.A = np.vstack([x, np.ones(len(x))]).T
    
    def reset(self):
        """
        Reset the calculator.
        """
        self.previous_tangent_force.clear()
        self.now_rate = 0.0
    
    def calculate_rate(self, tangent_force : float):
        """
        Calculate the rate of change of tangent force.
        Args:
            tangent_force (float): The current tangent force value.
        Returns:
            float: The rate of change of tangent force.
        """

        self.previous_tangent_force.append(tangent_force)
        if len(self.previous_tangent_force) < self.window_size:
            return 0.0
        
        # Calculate the rate of change with all the data in the window using linear fitting

        y = np.array(self.previous_tangent_force)

        m, c = np.linalg.lstsq(self.A, y, rcond=None)[0]
        self.now_rate = self.now_rate * self.momentum + m * (1 - self.momentum)
        return self.now_rate
        

exp_dir_list = []
for dir in os.listdir(data_dir):
    exp_dir = os.path.join(data_dir, dir)
    if not os.path.isdir(exp_dir):
        continue

    exp_dir_list.append(exp_dir)

# sort the directories by their names
exp_dir_list.sort(key=lambda x: int(os.path.basename(x).split('_')[1]), reverse=False)
rate_calculator = TangentForceRateCalculator(window_size=5, momentum=0.9)
for exp_dir in exp_dir_list:

    plt.figure(figsize=(10, 9))
    for sensor_id in [0,1]:
        tacdata, true_delta, predicted_delta, grasp_force, goal_force, better_delta , better_mu = load_data(exp_dir, sensor_id)
        if tacdata is None:
            continue

        sensored_force = np.sum(tacdata[:, :, 6:9], axis=1)
        tangent_force = np.linalg.norm(sensored_force[:,0:2], axis=1)
        normal_force = -sensored_force[:, 2]


        plt.subplot(3, 2, sensor_id + 1)
        plt.plot(tangent_force, label='Tangent Force')
        plt.plot(goal_force, label='Goal Force', linestyle='--')
        plt.plot(normal_force, label='Normal Force')
        plt.plot(grasp_force, label='Grasp Force')
        plt.ylabel('Force (N)')
        plt.legend()

        plt.subplot(3, 2, sensor_id + 3)
        plt.plot(true_delta, label='True Delta', color='blue', linewidth=2)
        plt.plot(predicted_delta, label='Estimated Delta', color='orange', linestyle='--', linewidth=2)
        plt.plot(better_delta, label='better Delta', color='green', linestyle=':')
        # plt.plot(better_mu, label='better Mu', color='red', linestyle='-.')
        # plt.plot(tangent_force/normal_force, label='Mu', color='blue', linestyle='-.')
        plt.title('Slip Delta Estimation')
        # plt.xlabel('Sample Index')
        plt.ylabel('Slip Delta')
        plt.legend()
        plt.subplot(3, 2, sensor_id + 5)
        contact_center_row, contact_center_col = calc_contact_center(tacdata)
        plt.plot(contact_center_row, label='Contact Center Row', color='green')
        plt.plot(contact_center_col, label='Contact Center Column', color='blue')
        plt.title('Contact Center')
        plt.xlabel('Sample Index')
        plt.ylabel('Contact Center Position')
        plt.legend()
        plt.tight_layout()

        mae = np.mean(np.abs(predicted_delta - true_delta))
        # print(f"Sensor {sensor_id} MAE: {mae:.4f}")
    print(exp_dir)
    plt.show()
